package com.wu.search.util;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.ObjectUtils;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.script.mustache.MultiSearchTemplateRequest;
import org.elasticsearch.script.mustache.MultiSearchTemplateResponse;
import org.elasticsearch.script.mustache.SearchTemplateRequest;
import org.elasticsearch.script.mustache.SearchTemplateResponse;

import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.stream.Collectors;

@Slf4j
public class RestHighLevelClientUtils {

    /**
     * 使用默认查询超时
     * @param restHighLevelClient
     * @param request
     * @return
     */
    public static SearchResponse getSearchResponse(RestHighLevelClient restHighLevelClient, SearchTemplateRequest request) {
        return getBaseSearchResponse(restHighLevelClient, request, null);
    }

    /**
     * 需要设置超时时间的调用
     * @param restHighLevelClient
     * @param request
     * @param timeOut
     * @return
     */
    public static SearchResponse getSearchResponse(RestHighLevelClient restHighLevelClient, SearchTemplateRequest request, Integer timeOut){
        return getBaseSearchResponse(restHighLevelClient, request, timeOut);
    }

    private static SearchResponse getBaseSearchResponse(RestHighLevelClient restHighLevelClient, SearchTemplateRequest request, Integer timeOut) {
        SearchTemplateResponse response = null;
        long startTime = System.currentTimeMillis();
        try {
            CompletableFuture<SearchTemplateResponse> future = CompletableFuture.supplyAsync(() -> {
                try {
                    return restHighLevelClient.searchTemplate(request, RequestOptions.DEFAULT);
                } catch (Exception e) {
                    log.error("======搜索 getBaseSearchResponse: e", e);
                }
                return null;
            });
            response = future.get(Optional.ofNullable(timeOut).orElse(500), TimeUnit.MILLISECONDS);
        } catch (InterruptedException | ExecutionException e) {
            log.error("======搜索 getBaseSearchResponse异常, params:{} ,reason:{}", request.getScriptParams().toString(), e.getMessage());
        } catch (TimeoutException e) {
            log.error("======搜索 getBaseSearchResponse超时, params:{}",request.getScriptParams().toString());
        }
        if (ObjectUtils.isNotEmpty(response)) {
            if (System.currentTimeMillis() - startTime > 200) {
                log.info("======搜索慢查询 getBaseSearchResponse params:{} es查询耗时:{} 总耗时：{}", request.getScriptParams(), response.getResponse().getTook().getMillis() , System.currentTimeMillis() - startTime);
            }
            return response.getResponse();
        }
        return null;
    }

    private static MultiSearchTemplateResponse getBaseMultiSearchResponse(RestHighLevelClient restHighLevelClient, MultiSearchTemplateRequest request, Integer timeOut) {
        MultiSearchTemplateResponse response = null;
        long startTime = System.currentTimeMillis();
        try {
            CompletableFuture<MultiSearchTemplateResponse> future = CompletableFuture.supplyAsync(() -> {
                try {
                    return restHighLevelClient.msearchTemplate(request, RequestOptions.DEFAULT);
                } catch (Exception e) {
                    log.error("======搜索 getBaseMultiSearchResponse: e", e);
                }
                return null;
            });
            response = future.get(Optional.ofNullable(timeOut).orElse(500), TimeUnit.MILLISECONDS);
        } catch (InterruptedException | ExecutionException e) {
            log.error("======搜索 getBaseMultiSearchResponse异常, params:{} ,reason:{}", request.requests().stream().map(SearchTemplateRequest::getScriptParams).map(Object::toString).collect(Collectors.joining(",")), e.getMessage());
        } catch (TimeoutException e) {
            log.error("======搜索 getBaseMultiSearchResponse超时, params:{}", request.requests().stream().map(SearchTemplateRequest::getScriptParams).map(Object::toString).collect(Collectors.joining(",")));
        }
        if (ObjectUtils.isNotEmpty(response)) {
            if (System.currentTimeMillis() - startTime > 200) {
                log.info("======搜索慢查询 getBaseMultiSearchResponse params:{} es查询耗时：{} 总耗时:{} ", request.requests().stream().map(SearchTemplateRequest::getScriptParams).map(Object::toString).collect(Collectors.joining(",")), response.getTook().getMillis(),System.currentTimeMillis() - startTime );
            }
            return response;
        }
        return null;
    }

    /**
     * 多请求异步查询
     * @param restHighLevelClient
     * @param multiRequest
     * @return
     */
    public static MultiSearchTemplateResponse getMultiSearchResponse(RestHighLevelClient restHighLevelClient, MultiSearchTemplateRequest multiRequest) {
        return getBaseMultiSearchResponse(restHighLevelClient, multiRequest, null);
    }

    public static MultiSearchTemplateResponse getMultiSuggestsResponse(RestHighLevelClient restHighLevelClient, MultiSearchTemplateRequest multiRequest){
        return getBaseMultiSuggestsResponse(restHighLevelClient, multiRequest, null);
    }

    public static MultiSearchTemplateResponse getBaseMultiSuggestsResponse(RestHighLevelClient restHighLevelClient, MultiSearchTemplateRequest request, Integer timeOut) {
        MultiSearchTemplateResponse response = null;
        long startTime = System.currentTimeMillis();
        try {
            CompletableFuture<MultiSearchTemplateResponse> future = CompletableFuture.supplyAsync(() -> {
                try {
                    return restHighLevelClient.msearchTemplate(request, RequestOptions.DEFAULT);
                } catch (Exception e) {
                    log.error("======搜索 getBaseMultiSuggestsResponse: e", e);
                }
                return null;
            });
            response = future.get(Optional.ofNullable(timeOut).orElse(500), TimeUnit.MILLISECONDS);
        } catch (InterruptedException | ExecutionException e) {
            log.error("======搜索 getBaseMultiSuggestsResponse异常, params:{} ,reason:{}", request.toString(), e.getMessage());
        } catch (TimeoutException e) {
            log.error("======搜索 getBaseMultiSuggestsResponse超时, params:{}", request.toString());
        }
        if (ObjectUtils.isNotEmpty(response)) {
            if (System.currentTimeMillis() - startTime > 200) {
                log.info("======搜索慢查询 getBaseMultiSuggestsResponse params:{} es查询耗时：{} 总耗时:{} " , request.toString(), response.getTook().getMillis(), System.currentTimeMillis() - startTime);
            }
            return response;
        }
        return null;
    }
}
